package xdoc;

import io.AceDocumentHandler;
import io.AdHocAceDocumentParser;
import io.AceOutputDocument;
import io.AceInputDocument;

import com.aliasi.chunk.Chunk;
import com.aliasi.chunk.Chunker;
import com.aliasi.chunk.Chunking;
import com.aliasi.cluster.ClusterScore;
import com.aliasi.dict.ExactDictionaryChunker;
import com.aliasi.hmm.HiddenMarkovModel;
import com.aliasi.hmm.HmmDecoder;
import com.aliasi.sentences.IndoEuropeanSentenceModel;
import com.aliasi.sentences.SentenceChunker;
import com.aliasi.sentences.SentenceModel;
import com.aliasi.tokenizer.IndoEuropeanTokenizerFactory;
import com.aliasi.tokenizer.TokenizerFactory;
import com.aliasi.tracker.Dictionary;
import com.aliasi.tracker.EntityPhraseChunker;
import com.aliasi.tracker.EntityUniverse;
import com.aliasi.tracker.InputDocument;
import com.aliasi.tracker.OutputDocument;
import com.aliasi.tracker.XDocCoref;
import com.aliasi.util.AbstractCommand;
import com.aliasi.util.AbstractExternalizable;
import com.aliasi.util.Files;
import com.aliasi.util.Streams;
import com.aliasi.util.Strings;
import com.aliasi.xml.DelegatingHandler;
import java.io.File;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.Set;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.SAXException;
import org.xml.sax.XMLReader;
import org.xml.sax.helpers.XMLReaderFactory;

public class RunXdoc extends AbstractCommand {

    public RunXdoc (String[] args) {
        super(args);
    }

    public void run() {
        HashMap<String,AceInputDocument> id2Document = new HashMap();
        HashMap<String,Set<XDocMention>> entityId2mentionResponse = new HashMap();
        try {

            File docsDir = getArgumentFile(DOCUMENTS_PARAM);
            File processedDocumentsDir =getArgumentFile(PROCESSED_DOCUMENTS_PARAM);
            File dictFile = getArgumentFile(DICT_PARAM);
            File neModelFile = getArgumentFile(NEMODEL_PARAM);
            File posModelFile = getArgumentFile(POSMODEL_PARAM);
            String tokenCountFileName = getArgument(TOKENS_PARAM);
            Double cutoff = null;

            if (hasArgument(CUTOFF_PARAM)) {
                cutoff = getArgumentDouble(CUTOFF_PARAM);
            }

            Dictionary dictionary
                = Dictionary.read(dictFile,com.aliasi.util.Strings.UTF8);

            TokenizerFactory tokenizerFactory
                = IndoEuropeanTokenizerFactory.FACTORY;
            
            SentenceModel sentenceModel
                = new IndoEuropeanSentenceModel(true,false);
            
            SentenceChunker sentenceChunker
                = new SentenceChunker(tokenizerFactory,sentenceModel);

            Chunker neChunker
                = (Chunker) AbstractExternalizable.readObject(neModelFile);

            HiddenMarkovModel posHmm
                = (HiddenMarkovModel) AbstractExternalizable.readObject(posModelFile);

            HmmDecoder posTagger= new HmmDecoder(posHmm);
            
            String phraseCounts 
                = Files.readFromFile(new File(tokenCountFileName));
            
            String[] lines = phraseCounts.split("\n");
            
            ExactDictionaryChunker dictChunker 
                = dictionary.chunker(tokenizerFactory,
                                     RETURN_ALL_MATCHES_FALSE,
                                     CASE_SENSITIVE_FALSE);

            EntityPhraseChunker entityPhraseChunker 
                = new AceEntityPhraseChunker(neChunker,
                                             posTagger,
                                             dictChunker,
                                             dictionary.stopPhrases(),
                                             STOP_SUBSTRING_LIST,
                                             new HashMap <String, Integer>(),
                                             cutoff);

            AceTracker tracker = new AceTracker(tokenizerFactory,
                                          sentenceChunker,
                                          entityPhraseChunker,
                                          dictionary);
            parseAdhocDocs(docsDir,id2Document,SGML_SUFFIX);

            runDocs(tracker, id2Document, processedDocumentsDir, entityId2mentionResponse);

            //System.out.println("-------------------Response-------------");
            //printAnnotations(entityId2mentionResponse);
            
              XDocCoref xdoc = tracker.xDocCoref();
                System.out.println(xdoc.entityUniverse());

        }
        catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    void parseAdhocDocs(File fileOrDir, HashMap<String,AceInputDocument> id2Document, String fileFilter) 
        throws IOException{
        if (fileOrDir.isDirectory()) {
            File[] files = fileOrDir.listFiles();
            for (int i = 0; i < files.length; ++i) {
                parseAdhocDocs(files[i],id2Document,fileFilter);
            }
        } 
        else if (fileOrDir.toString().endsWith(fileFilter)) {
            {
                AceInputDocument doc = AdHocAceDocumentParser.parse(Files.readFromFile(fileOrDir));
                id2Document.put(doc.id(),doc);
            }
        }
    }

    void parseDocs(File fileOrDir, DelegatingHandler docHandler, String fileFilter) 
        throws IOException, SAXException {
        if (fileOrDir.isDirectory()) {
            File[] files = fileOrDir.listFiles();
            for (int i = 0; i < files.length; ++i) {
                parseDocs(files[i],docHandler,fileFilter);
            }
        } 
        else if (fileOrDir.toString().endsWith(fileFilter)) {
            {
                loadXML(fileOrDir, docHandler);
            }
        }
    }
   
    void runDocs (AceTracker tracker, HashMap<String,AceInputDocument> id2Document,
                  File annotationsDir, Map<String,Set<XDocMention>> entityId2mentionResponse) 
        throws IOException {
        Set<Map.Entry<String,AceInputDocument>> entrySet = id2Document.entrySet();
        for (Map.Entry<String,AceInputDocument> entry  : entrySet) {
            //Map map = tracker.processDocumentNoOutput(entry.getValue();
            AceInputDocument doc = entry.getValue();
            if (entry.getValue().isUnProcessableDocument()) {
                System.out.println("****************doc is not processable" +  doc);
                continue;
            }
            //            System.out.println("Doc is: " + doc);
            
            //AceOutputDocument annotation = tracker.processDocument(doc,entityId2mentionResponse);
            OutputDocument annotation = tracker.processDocumentOffsets(doc);
            Files.writeStringToFile(annotation.toString(), new File(annotationsDir,annotation.id() + ".apf.xml"));
        }
     }

    void printDocs (HashMap<String,InputDocument> id2Document ) {
        Set<Map.Entry<String,InputDocument>> entrySet = id2Document.entrySet();
        for (Map.Entry<String,InputDocument> entry  : entrySet) {
            System.out.print(entry.getKey() + "\n" + entry.getValue());
        }
    }

    void printAnnotations(HashMap<String,Set<XDocMention>> entityId2mention) {
        Set<Map.Entry<String,Set<XDocMention>>> entrySet = entityId2mention.entrySet();
        for (Map.Entry<String,Set<XDocMention>> entry  : entrySet) {
            System.out.println(entry.getKey());
            for (XDocMention ment : entry.getValue()) {
                System.out.print("mention:\t" + ment.mention() + "\n");
            }
        }
    }

    void printAnnotationsNoteOverlaps(HashMap<String,Set<String>> entityId2mention) {
        Set seen = new HashSet();
        Set<Map.Entry<String,Set<String>>> entrySet = entityId2mention.entrySet();
        for (Map.Entry<String,Set<String>> entry  : entrySet) {
            System.out.println(entry.getKey());
            for (String str : entry.getValue()) {
                if (seen.contains(str)) {
                    System.out.println("seen before: " + str);
                }
                seen.add(str);
                System.out.print("\t" + str + "\n");
            }
        }
    }

    
    static void loadXML(File input, DelegatingHandler handler) 
        throws SAXException, IOException {
        //   System.out.println("parsing " + input);
        XMLReader saxParser = XMLReaderFactory.createXMLReader();
        saxParser.setContentHandler(handler);
        String url = Files.fileToURLName(input);
        InputSource in = new InputSource(url);
        saxParser.parse(in);
        //        System.out.println("Loaded " + input);
    }

    private static final String[] STOP_SUBSTRING_LIST = {"to"};

    private static final boolean RETURN_ALL_MATCHES_FALSE = false;
    private static final boolean CASE_SENSITIVE_FALSE = false;
    private static final String DOCUMENTS_PARAM = "documents";
    private static final String PROCESSED_DOCUMENTS_PARAM = "processedDocuments";
    private static final String DICT_PARAM = "dictionary";
    private static final String NEMODEL_PARAM = "neModel";
    private static final String POSMODEL_PARAM = "posModel";
    private static final String TOKENS_PARAM = "tokens";
    private static final String ANNOTATIONS_PARAM = "annotations";
    private static final String XML_HEADER = "<?xml version=\"1.0\"?>";
    private static final String APF_XML_SUFFIX = "apf.xml";
    private static final String SGML_SUFFIX = "sgm";
    private static final String CUTOFF_PARAM = "cutoff";

    public static void main (String[] args) {
        RunXdoc cmd = new RunXdoc(args);
        cmd.run();
    }
}
